import os
import typing as t
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager

import rich.live
import rich.progress
import rich.style
import rich.table


def index_column():
    return rich.progress.TextColumn(
        '{task.fields[index]}',
        style=rich.style.Style(
            color='bright_black',
        ),
        justify='right',
    )


class Progress:
    def __init__(
        self,
        console: rich.console.Console,
    ):
        self.root = rich.table.Table.grid()
        self.current = rich.progress.Progress(
            # index_column(),
            rich.progress.SpinnerColumn(style='yellow'),
            rich.progress.TextColumn('{task.description}'),
        )
        self.total = rich.progress.Progress(
            rich.progress.TextColumn(
                '[progress.description]{task.description} [progress.filesize][{task.completed}/{task.total}]'
            ),
            rich.progress.BarColumn(),
            rich.progress.TaskProgressColumn(),
            rich.progress.TimeRemainingColumn(),
        )
        self.root.add_row(self.current)
        self.root.add_row(self.total)
        self.live = LiveProgress(
            self.root,
            auto_refresh=True,
            refresh_per_second=1,
            console=console,
            transient=True,
        )
        self.executor = ThreadPoolExecutor(max_workers=os.cpu_count())


@contextmanager
def progress_executor(console: rich.console.Console):
    p = Progress(console)
    with p.live:
        with p.executor:
            yield p


PTask = t.TypeVar('PTask')


def progress_run(
    console: rich.console.Console,
    tasks: list[PTask],
    func: t.Callable[[PTask], tuple | None],
):
    p: Progress
    try:
        with progress_executor(console) as p:
            total_id = p.total.add_task('Total', start=True, total=len(tasks))
            p.live.refresh()

            def task_run(task: PTask, index: int = 0):
                task_id = p.current.add_task(str(task), index=index + 1, total=None)
                p.live.refresh()
                try:
                    result = func(task)
                    if result is not None:
                        p.live.log(*result)
                except Exception:
                    p.live.print_exception(show_locals=True)
                p.total.advance(total_id)
                p.current.remove_task(task_id)
                p.live.refresh()

            for index, task in enumerate(tasks):
                p.executor.submit(task_run, task, index)
    except KeyboardInterrupt:
        p.executor.shutdown(wait=False, cancel_futures=True)
        raise


class LiveProgress(rich.live.Live):
    def __init__(
        self,
        renderable=None,
        *,
        console=None,
        screen=False,
        auto_refresh=True,
        refresh_per_second=4,
        transient=False,
        redirect_stdout=True,
        redirect_stderr=True,
        vertical_overflow='ellipsis',
        get_renderable=None,
    ):
        super().__init__(
            renderable,
            console=console,
            screen=screen,
            auto_refresh=auto_refresh,
            refresh_per_second=refresh_per_second,
            transient=transient,
            redirect_stdout=redirect_stdout,
            redirect_stderr=redirect_stderr,
            vertical_overflow=vertical_overflow,
            get_renderable=get_renderable,
        )

    def log(self, *args, **kwargs) -> None:
        """Log rich content to the terminal.

        Args:
            objects (positional args): Objects to log to the terminal.
            sep (str, optional): String to write between print data. Defaults to " ".
            end (str, optional): String to write at end of print data. Defaults to "\\\\n".
            style (Union[str, Style], optional): A style to apply to output. Defaults to None.
            justify (str, optional): One of "left", "right", "center", or "full". Defaults to ``None``.
            emoji (Optional[bool], optional): Enable emoji code, or ``None`` to use console default. Defaults to None.
            markup (Optional[bool], optional): Enable markup, or ``None`` to use console default. Defaults to None.
            highlight (Optional[bool], optional): Enable automatic highlighting, or ``None`` to use console default. Defaults to None.
            log_locals (bool, optional): Boolean to enable logging of locals where ``log()``
                was called. Defaults to False.
            _stack_offset (int, optional): Offset of caller from end of call stack. Defaults to 1.
        """
        with self._lock:
            self.console.log(*args, **kwargs)

    def print_exception(self, **kwargs) -> None:
        with self._lock:
            self.console.print_exception(**kwargs)

    def __enter__(self) -> t.Self:
        super().__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        super().__exit__(exc_type, exc_val, exc_tb)
